#include <windows.h>
#include <shlobj.h>
#include <shobjidl.h>
#include <objbase.h>
#include <commctrl.h> // 包含列表视图控件的消息定义，如LVM_GETITEMCOUNT

class DesktopIconOrganizer
{
public:
    DesktopIconOrganizer() : pShellView(nullptr), pDesktopFolder(nullptr) {}

    bool initialize()
    {
        // 使用COINIT_APARTMENTTHREADED，因为Shell API通常需要单线程单元
        HRESULT hr = CoInitializeEx(NULL, COINIT_APARTMENTTHREADED);
        if (FAILED(hr)) {
            MessageBox(NULL, L"COM初始化失败", L"错误", MB_ICONERROR);
            return false;
        }

        // 获取桌面文件夹
        hr = SHGetDesktopFolder(&pDesktopFolder);
        if (FAILED(hr)) {
            MessageBox(NULL, L"获取桌面文件夹失败", L"错误", MB_ICONERROR);
            CoUninitialize();
            return false;
        }

        // 获取桌面Shell视图
        hr = pDesktopFolder->CreateViewObject(NULL, IID_IShellView, (void**)&pShellView);
        if (FAILED(hr)) {
            MessageBox(NULL, L"创建Shell视图失败", L"错误", MB_ICONERROR);
            pDesktopFolder->Release();
            CoUninitialize();
            return false;
        }

        return true;
    }

    void cleanup()
    {
        if (pShellView) {
            pShellView->Release();
            pShellView = nullptr;
        }
        if (pDesktopFolder) {
            pDesktopFolder->Release();
            pDesktopFolder = nullptr;
        }
        CoUninitialize();
    }

    // 将图标整理到指定位置的网格中
    bool arrangeIconsToGrid(int startX, int startY, int horizontalSpacing = 100, int verticalSpacing = 100)
    {
        if (!pShellView) return false;

        HWND hwndListView;
        HRESULT hr = pShellView->GetWindow(&hwndListView);
        if (FAILED(hr) || !hwndListView) return false;

        // 获取桌面图标数量
        int itemCount = (int)SendMessage(hwndListView, LVM_GETITEMCOUNT, 0, 0);
        if (itemCount <= 0) return false;

        // 获取桌面工作区大小，避免图标摆放到任务栏区域
        RECT workArea;
        SystemParametersInfo(SPI_GETWORKAREA, 0, &workArea, 0);
        int workAreaWidth = workArea.right - workArea.left;
        int workAreaHeight = workArea.bottom - workArea.top;

        // 计算每行可容纳的图标数量
        int iconsPerRow = (workAreaWidth - startX) / horizontalSpacing;
        if (iconsPerRow <= 0) iconsPerRow = 1;

        // 设置每个图标的位置
        for (int i = 0; i < itemCount; i++) {
            int row = i / iconsPerRow;
            int col = i % iconsPerRow;

            int posX = startX + col * horizontalSpacing;
            int posY = startY + row * verticalSpacing;

            // 确保图标在工作区内
            if (posX >= workAreaWidth) posX = workAreaWidth - horizontalSpacing;
            if (posY >= workAreaHeight) posY = workAreaHeight - verticalSpacing;

            // 设置图标位置
            SendMessage(hwndListView, LVM_SETITEMPOSITION, i, MAKELPARAM(posX, posY));
        }

        // 刷新桌面
        refreshDesktop();
        return true;
    }

private:
    IShellFolder* pDesktopFolder;
    IShellView* pShellView;

    void refreshDesktop()
    {
        // 方法1: 通过查找桌面窗口并发送刷新消息
        HWND hDesktop = FindWindow(L"Progman", L"Program Manager");
        if (hDesktop) {
            // 发送刷新消息，0x740C是WM_USER+0x102C，这是一个已知的桌面刷新消息
            SendMessageTimeout(hDesktop, 0x740C, 0, 0, SMTO_NORMAL, 1000, NULL);
        }

        // 方法2: 强制重绘桌面
        HWND hShellWindow = GetShellWindow();
        if (hShellWindow) {
            InvalidateRect(hShellWindow, NULL, TRUE);
            UpdateWindow(hShellWindow);
        }
    }
};